#include "Pass.h"
#include "ComSubExprEli.h"
#include <set>
#include <map>
#include <vector>
#include <algorithm>
using namespace std;

std::set<Instruction *, cmp_inst> all_expr;                       //所有的可用表达式
std::map<BasicBlock *, std::set<Instruction *, cmp_inst>> bb_in;  //块B入口的可用表达式集合
std::map<BasicBlock *, std::set<Instruction *, cmp_inst>> bb_out; //块B出口的可用表达式集合
std::map<BasicBlock *, std::set<Instruction *, cmp_inst>> bb_gen; //块B产生的可用表达式集合

std::set<Instruction *, cmp_inst> Set_intersection(std::set<Instruction *, cmp_inst> &a, std::set<Instruction *, cmp_inst> &b)
{
    std::set<Instruction *, cmp_inst> ret = {};
    std::set_intersection(a.begin(), a.end(), b.begin(), b.end(), inserter(ret, ret.begin()));
    return ret;
}

std::set<Instruction *, cmp_inst> Set_union(std::set<Instruction *, cmp_inst> &a, std::set<Instruction *, cmp_inst> &b)
{
    std::set<Instruction *, cmp_inst> ret = {};
    std::set_union(a.begin(), a.end(), b.begin(), b.end(), inserter(ret, ret.begin()));
    return ret;
}

bool is_same_set(std::set<Instruction *, cmp_inst> &a, std::set<Instruction *, cmp_inst> &b)
{
    if (a.size() != b.size())
        return false;
    else
        return true;
    // std::map<std::set<Instruction *, cmp_inst>, int> mTemp; //借助于map的数据结构判断new_bb_out和旧的bb_out两个集合是否相等
    // mTemp[a] = 1;
    // if (mTemp.find(b) != mTemp.end())
    // {
    //     return false;
    // }
    // else
    // {
    //     return true;
    // }
}

void ComSubExprEli::compute_gen(Function *fun)
{
    std::list<Instruction *> delete_inst;
    for (auto bb : fun->get_basic_blocks()) //实现局部的公共子表达式删除和公共子表达式gen的计算，由于这是基于ssa格式的代码，故每个块的公共子表达式的kill都是空集，不需要计算
    {
        for (auto inst : bb->get_instructions())
        {
            if (is_valid_expr(inst))
            {
                auto flag = bb_gen[bb].insert(inst);
                if (flag.second == false) //插入失败，在基本块内之前就已经有过该表达式了
                {
                    auto last_inst = bb_gen[bb].find(inst);
                    inst->replace_all_use_with(*last_inst);
                    delete_inst.push_back(inst);
                }
                else
                {
                    all_expr.insert(inst);
                }
            }
        }
        for (auto inst : delete_inst)
            bb->delete_instr(inst);
    }
}

void ComSubExprEli::available_expr_analysis(Function *fun) //计算每个基本块的in和out
{
    for (auto bb : fun->get_basic_blocks()) //由于合流算符是交，所有块的out初始化为所有可用的公共子表达式
    {
        bb_out[bb] = all_expr;
    }
    bb_out[fun->get_entry_block()] = std::set<Instruction *, cmp_inst>(); //entry的out应当设为空集
    bool flag = true;

    while (flag)
    {
        flag = false;
        for (auto bb : fun->get_basic_blocks())
        {
            if (bb != fun->get_entry_block())
            {
                std::set<Instruction *, cmp_inst> temp_in = {};
                if (bb->get_pre_basic_blocks().size() != 0)
                {
                    temp_in = bb_out[bb->get_pre_basic_blocks().front()];
                    for (auto pre_bb : bb->get_pre_basic_blocks()) //bb_in应当是所有前驱块的交集
                    {
                        temp_in = Set_intersection(temp_in, bb_out[pre_bb]);
                    }
                }
                bb_in[bb] = temp_in;
                std::set<Instruction *, cmp_inst> new_bb_out = Set_union(bb_in[bb], bb_gen[bb]);
                if (!is_same_set(new_bb_out, bb_out[bb]))
                {
                    flag = true;
                }
                bb_out[bb]=new_bb_out;
            }
        }
    }
}

void ComSubExprEli::global_common_expr_delete(Function *fun)
{ //由于全局的公共子表达式可能存在嵌套的情况，故不能按照块内公共子表达式的方法来直接删除，必须溯源use-def链，找到最早的公共子表达式
    std::set<Instruction *> repetitive_inst = {};
    std::map<Instruction *, Instruction *> old_inst_map;
    for (auto bb : fun->get_basic_blocks())
    {
        for (auto inst : bb->get_instructions())
        {
            if (is_valid_expr(inst))
            {
                auto old_inst = bb_in[bb].find(inst);
                if (old_inst != bb_in[bb].end()) //找到一条公共子表达式
                {
                    if (*old_inst != inst)
                    {
                        old_inst_map[inst] = *old_inst;
                        repetitive_inst.insert(inst);
                    }
                }
            }
        }
    }
    for (auto inst : repetitive_inst)
    {
        Instruction *temp_inst = nullptr;
        for (temp_inst = old_inst_map[inst]; old_inst_map.find(temp_inst) != old_inst_map.end(); temp_inst = old_inst_map[temp_inst])
            ;
        inst->replace_all_use_with(temp_inst);
        inst->get_parent()->delete_instr(inst);
    }
}

void ComSubExprEli::execute()
{
    for (auto fun : module->get_functions())
    {
        for (auto bb : fun->get_basic_blocks())
        {
            bb_in[bb] = std::set<Instruction *, cmp_inst>();
            bb_out[bb] = std::set<Instruction *, cmp_inst>();
            bb_gen[bb] = std::set<Instruction *, cmp_inst>();
        }
        compute_gen(fun);
        available_expr_analysis(fun);
        // global_common_expr_delete(fun);
    }
    module->set_print_name();
    /*you need to finish this function*/
}

bool ComSubExprEli::is_valid_expr(Instruction *inst)
{
    return !(inst->is_void() || inst->is_call() || inst->is_phi() || inst->is_alloca() || inst->is_load() || inst->is_cmp() || inst->is_zext()); //TODO:CHECK VALID INST
}